import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw
from skimage.transform import resize
import random

def create_directory(directory):
    """Create directory if it doesn't exist."""
    if not os.path.exists(directory):
        os.makedirs(directory)

def generate_star_field(width, height, num_stars=100):
    """Generate a star field image."""
    image = np.zeros((height, width))
    
    # Add stars (bright pixels)
    for _ in range(num_stars):
        x = np.random.randint(0, width)
        y = np.random.randint(0, height)
        brightness = np.random.uniform(0.3, 1.0)
        image[y, x] = brightness
        
        # Sometimes add a small glow around stars
        if random.random() > 0.7:
            size = random.randint(1, 3)
            y_min = max(0, y - size)
            y_max = min(height, y + size + 1)
            x_min = max(0, x - size)
            x_max = min(width, x + size + 1)
            
            for i in range(y_min, y_max):
                for j in range(x_min, x_max):
                    if (i != y or j != x) and random.random() > 0.5:
                        image[i, j] = brightness * random.uniform(0.1, 0.4)
    
    return image

def add_asteroid(image, min_size=3, max_size=8):
    """Add an asteroid to the image."""
    height, width = image.shape
    
    # Asteroid properties
    asteroid_size = random.randint(min_size, max_size)
    x = random.randint(asteroid_size, width - asteroid_size)
    y = random.randint(asteroid_size, height - asteroid_size)
    
    # Create asteroid shape (irregular circle)
    img = Image.fromarray(np.uint8(image * 255))
    draw = ImageDraw.Draw(img)
    
    # Create irregular shape for asteroid
    points = []
    for angle in range(0, 360, 30):
        radius = asteroid_size * random.uniform(0.7, 1.0)
        px = x + radius * np.cos(np.radians(angle))
        py = y + radius * np.sin(np.radians(angle))
        points.append((px, py))
    
    # Draw the asteroid
    draw.polygon(points, fill=255)
    
    # Convert back to numpy array
    result = np.array(img) / 255.0
    
    return result, (x, y, asteroid_size)

def generate_dataset(num_images=100, image_size=(64, 64), with_asteroid_prob=0.5, output_dir="dataset"):
    """Generate a dataset of images with and without asteroids."""
    create_directory(output_dir)
    create_directory(os.path.join(output_dir, "with_asteroid"))
    create_directory(os.path.join(output_dir, "without_asteroid"))
    
    metadata = []
    
    for i in range(num_images):
        # Generate star field
        image = generate_star_field(image_size[0], image_size[1], num_stars=random.randint(50, 150))
        
        # Decide whether to add asteroid
        has_asteroid = random.random() < with_asteroid_prob
        asteroid_info = None
        
        if has_asteroid:
            image, asteroid_info = add_asteroid(image)
            save_path = os.path.join(output_dir, "with_asteroid", f"image_{i}.png")
        else:
            save_path = os.path.join(output_dir, "without_asteroid", f"image_{i}.png")
        
        # Save image
        plt.imsave(save_path, image, cmap='gray')
        
        # Add to metadata
        metadata.append({
            "image_id": i,
            "has_asteroid": has_asteroid,
            "asteroid_info": asteroid_info,
            "path": save_path
        })
        
        # Display progress
        if (i + 1) % 10 == 0:
            print(f"Generated {i + 1}/{num_images} images")
    
    return metadata

if __name__ == "__main__":
    print("Generating dataset of space images with and without asteroids...")
    metadata = generate_dataset(num_images=200, image_size=(64, 64))
    print("Dataset generation complete!")
    
    # Display some examples
    fig, axes = plt.subplots(2, 3, figsize=(12, 8))
    
    # Images with asteroids
    with_asteroid = [m for m in metadata if m["has_asteroid"]]
    for i in range(3):
        if i < len(with_asteroid):
            img = plt.imread(with_asteroid[i]["path"])
            axes[0, i].imshow(img, cmap='gray')
            axes[0, i].set_title(f"With Asteroid")
            axes[0, i].axis('off')
    
    # Images without asteroids
    without_asteroid = [m for m in metadata if not m["has_asteroid"]]
    for i in range(3):
        if i < len(without_asteroid):
            img = plt.imread(without_asteroid[i]["path"])
            axes[1, i].imshow(img, cmap='gray')
            axes[1, i].set_title(f"Without Asteroid")
            axes[1, i].axis('off')
    
    plt.tight_layout()
    plt.savefig("sample_images.png")
    plt.show()
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw
from skimage.transform import resize
import random

def create_directory(directory):
    """Create directory if it doesn't exist."""
    if not os.path.exists(directory):
        os.makedirs(directory)

def generate_star_field(width, height, num_stars=100):
    """Generate a star field image."""
    image = np.zeros((height, width))
    
    # Add stars (bright pixels)
    for _ in range(num_stars):
        x = np.random.randint(0, width)
        y = np.random.randint(0, height)
        brightness = np.random.uniform(0.3, 1.0)
        image[y, x] = brightness
        
        # Sometimes add a small glow around stars
        if random.random() > 0.7:
            size = random.randint(1, 3)
            y_min = max(0, y - size)
            y_max = min(height, y + size + 1)
            x_min = max(0, x - size)
            x_max = min(width, x + size + 1)
            
            for i in range(y_min, y_max):
                for j in range(x_min, x_max):
                    if (i != y or j != x) and random.random() > 0.5:
                        image[i, j] = brightness * random.uniform(0.1, 0.4)
    
    return image

def add_asteroid(image, min_size=3, max_size=8):
    """Add an asteroid to the image."""
    height, width = image.shape
    
    # Asteroid properties
    asteroid_size = random.randint(min_size, max_size)
    x = random.randint(asteroid_size, width - asteroid_size)
    y = random.randint(asteroid_size, height - asteroid_size)
    
    # Create asteroid shape (irregular circle)
    img = Image.fromarray(np.uint8(image * 255))
    draw = ImageDraw.Draw(img)
    
    # Create irregular shape for asteroid
    points = []
    for angle in range(0, 360, 30):
        radius = asteroid_size * random.uniform(0.7, 1.0)
        px = x + radius * np.cos(np.radians(angle))
        py = y + radius * np.sin(np.radians(angle))
        points.append((px, py))
    
    # Draw the asteroid
    draw.polygon(points, fill=255)
    
    # Convert back to numpy array
    result = np.array(img) / 255.0
    
    return result, (x, y, asteroid_size)

def generate_dataset(num_images=100, image_size=(64, 64), with_asteroid_prob=0.5, output_dir="dataset"):
    """Generate a dataset of images with and without asteroids."""
    create_directory(output_dir)
    create_directory(os.path.join(output_dir, "with_asteroid"))
    create_directory(os.path.join(output_dir, "without_asteroid"))
    
    metadata = []
    
    for i in range(num_images):
        # Generate star field
        image = generate_star_field(image_size[0], image_size[1], num_stars=random.randint(50, 150))
        
        # Decide whether to add asteroid
        has_asteroid = random.random() < with_asteroid_prob
        asteroid_info = None
        
        if has_asteroid:
            image, asteroid_info = add_asteroid(image)
            save_path = os.path.join(output_dir, "with_asteroid", f"image_{i}.png")
        else:
            save_path = os.path.join(output_dir, "without_asteroid", f"image_{i}.png")
        
        # Save image
        plt.imsave(save_path, image, cmap='gray')
        
        # Add to metadata
        metadata.append({
            "image_id": i,
            "has_asteroid": has_asteroid,
            "asteroid_info": asteroid_info,
            "path": save_path
        })
        
        # Display progress
        if (i + 1) % 10 == 0:
            print(f"Generated {i + 1}/{num_images} images")
    
    return metadata

if __name__ == "__main__":
    print("Generating dataset of space images with and without asteroids...")
    metadata = generate_dataset(num_images=200, image_size=(64, 64))
    print("Dataset generation complete!")
    
    # Display some examples
    fig, axes = plt.subplots(2, 3, figsize=(12, 8))
    
    # Images with asteroids
    with_asteroid = [m for m in metadata if m["has_asteroid"]]
    for i in range(3):
        if i < len(with_asteroid):
            img = plt.imread(with_asteroid[i]["path"])
            axes[0, i].imshow(img, cmap='gray')
            axes[0, i].set_title(f"With Asteroid")
            axes[0, i].axis('off')
    
    # Images without asteroids
    without_asteroid = [m for m in metadata if not m["has_asteroid"]]
    for i in range(3):
        if i < len(without_asteroid):
            img = plt.imread(without_asteroid[i]["path"])
            axes[1, i].imshow(img, cmap='gray')
            axes[1, i].set_title(f"Without Asteroid")
            axes[1, i].axis('off')
    
    plt.tight_layout()
    plt.savefig("sample_images.png")
    plt.show()